import matplotlib.pyplot as plt
import numpy as np

import torch
import torchvision
import torchvision.transforms as transforms

import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

# transforms
transform = transforms.Compose(
    [transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))])

data_folder = 'E:/study_code/torch_study/aigc/data'

# datasets
trainset = torchvision.datasets.FashionMNIST(data_folder,
    download=True,
    train=True,
    transform=transform)
testset = torchvision.datasets.FashionMNIST(data_folder,
    download=True,
    train=False,
    transform=transform)

# dataloaders
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                        shuffle=True, num_workers=2)


testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                        shuffle=False, num_workers=2)

from torch.utils.tensorboard import SummaryWriter

# 设定工作记录档目录
writer = SummaryWriter('runs/fashion_mnist_experiment_1')
# 读取资料
# dataiter = iter(trainloader)
# images, labels = next(dataiter)
#
# # 建立图像方格
# img_grid = torchvision.utils.make_grid(images)
#
# # 写入 tensorboard
# writer.add_image('four_fashion_mnist_images', img_grid)